/*
    This file is part of the Time-of-Flight Tracer program. It is not part of
    the original PBRT source distribution. See the included license file for
    more information.

    Copyright(c) 2014 Microsoft Corporation

    Author: Phil Pitts
*/


#include "stdafx.h"
#include "integrators/tofpath.h"
#include "scene.h"
#include "intersection.h"
#include "paramset.h"


ToFPathIntegrator::ToFPathIntegrator(int md, float lambdaMin, float lambdaMax) :
PathIntegrator(md), lambdaMin(lambdaMin), lambdaMax(lambdaMax) {
}


IntegrationResult ToFPathIntegrator::Li(const Scene *scene,
	const Renderer *renderer, const RayDifferential &ray, 
	const Intersection &isect, const Sample *sample, RNG &rng, 
	MemoryArena &arena) const {
	// Declare common path integration variables
	Spectrum L(0.f), escapedAlpha(0.f);
	RayDifferential escapedRay;
	bool prevSpecular = true;
	PathData data(scene, renderer, rng, arena);
	ToFPathVertex *path = new ToFPathVertex[maxDepth];

	int pathLength = GeneratePath(data, ray, isect, sample, path, &escapedRay,
		&escapedAlpha);

	vector<HistogramSample> histSamples(pathLength);
	for (int i = 0; i < pathLength; i++) {
		Spectrum localL(0.f);

		// Initialize basic variables for camera path vertex
		const ToFPathVertex &v = path[i];

		// Add emitted light from vertex if appropriate
		if (prevSpecular)
			localL += v.alpha * v.isect.Le(v.wPrev);

		// Add direct illumination for path vertex
		float light2isect;
		localL += DirectLightingVertex(data, sample, i, v, &light2isect);

		prevSpecular = v.specularBounce;

		histSamples[i].distance = v.pathlength + light2isect;
		// TODO remove lambdas from Camera
		histSamples[i].radiance = localL.Average(lambdaMin, lambdaMax);
		L += localL;
	}

	delete[] path;
	return IntegrationResult(L, histSamples);
}


uint32_t ToFPathIntegrator::GeneratePath(PathData &data, const RayDifferential r,
	const Intersection &firstIsect, const Sample* sample, ToFPathVertex *path,
	RayDifferential *escapedRay, Spectrum *escapedAlpha) const {
	// Declare common path integration variables
	Spectrum pathThroughput = 1.f;
	RayDifferential ray(r);

	ToFPathVertex *v = &path[0];
	v->isect = firstIsect;

	int length = 0;
	for (; length < maxDepth; length++) {
		// Try to generate next vertex of ray path
		v = &path[length];
		if (length > 0) {
			if (!data.scene->Intersect(ray, &v->isect)) {
				*escapedAlpha = pathThroughput;
				*escapedRay = ray;
				break;
			}
		}
		SavePathVertexInfo(v, ray, pathThroughput, data.arena);
		SaveToFVertexInfo(path, length, ray);

		// Sample BSDF to get new path direction
		float pdf;
		BxDFType flags;
		BSDFSample BSDFSampleOut = BSDFSample(sample, pathSampleOffset, length);
		Spectrum f = v->bsdf->Sample_f(-ray.d, &v->wNext, BSDFSampleOut, &pdf,
			BSDF_ALL, &flags);
		v->specularBounce = (flags & BSDF_SPECULAR) != 0;
		v->nSpecularComponents = v->bsdf->NumComponents(BxDFType(BSDF_SPECULAR |
			BSDF_REFLECTION | BSDF_TRANSMISSION));

		if (f.IsBlack() || pdf == 0. || !RussianRoulette(data, *v, f, pdf, &ray,
			&pathThroughput)) {
			return length + 1;
		}
	}
	return length;
}


void ToFPathIntegrator::SaveToFVertexInfo(ToFPathVertex *path, int pathIdx, 
	const RayDifferential &ray) const {
	path[pathIdx].pathlength = (ray.o - path[pathIdx].isect.dg.p).Length();
	
	if (pathIdx > 0) {
		path[pathIdx].pathlength += path[pathIdx - 1].pathlength;
	}
}


Spectrum ToFPathIntegrator::DirectLightingVertex(PathData &data, 
	const Sample *sample, int pathIdx, const ToFPathVertex &v, 
	float *light2isect) const {
	Spectrum L(0.f);
	const Point &pc = v.bsdf->dgShading.p;
	const Normal &n = v.bsdf->dgShading.nn;

	float lightPdf;
	uint32_t lightNum = lightDistribution->SampleDiscrete(
		sample->oneD[lightNumOffset][pathIdx], &lightPdf);
	const Light *light = data.scene->lights[lightNum];
	BSDFSample bsdfSample(sample, bsdfSampleOffset, pathIdx);
	LightSample lightSample(sample, lightSampleOffset, pathIdx);

	L = v.alpha *
		EstimateDirect(data.scene, data.renderer, data.arena, light, pc, n, 
		v.wPrev, v.isect.rayEpsilon, sample->time, v.bsdf, data.rng, lightSample,
		bsdfSample, BxDFType(BSDF_ALL & ~BSDF_SPECULAR)) / lightPdf;

	// Extract distance from point to light for ToF histogram
	Normal normal;
	Ray lightRay;
	light->Sample_L(data.scene, lightSample, data.rng.RandomFloat(),
		data.rng.RandomFloat(), sample->time, &lightRay,
		&normal, &lightPdf);
	Point pl = lightRay.o;
	*light2isect = Distance(pl, pc);

	return L;
}


ToFPathIntegrator *CreateToFPathSurfaceIntegrator(
	const ParamSet &params) {
    int maxDepth = params.FindOneInt("maxdepth", 5);
	float lambdaMin = params.FindOneFloat("lambdamin", 400);
	float lambdaMax = params.FindOneFloat("lambdamax", 1000);
	return new ToFPathIntegrator(maxDepth, lambdaMin, lambdaMax);
}
